package com.taglab.chic.nonblocking;

import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.HashMap;
import java.util.Date;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 *
 */
public class NonBlockingCachingMethodInterceptor implements MethodInterceptor {

  private final Logger logger = LoggerFactory.getLogger( this.getClass() );

  /**
   * The number of seconds to cache method values for, before updating.
   */
  private int cacheTimeSeconds;

  /**
   * The number of milliseconds that threads should wait for an initial
   * cache to be created by another thread, before giving up and calling
   * on the underlying method themselves.
   */
  private int firstCacheWaitTimeMillis;

  /**
   * The number of milliseconds that threads should wait for an expired
   * cache to be updated by another thread, before giving up and returning
   * the old stale cache value.
   */
  private int updatedCacheWaitTimeMillis;

  /**
   * The reason that this isn't a ConcurrentMap is because we are separately
   * ensuring that only one thread can write to it at a time (except in the
   * case where we are caching the value for the first time and the first
   * thread takes too long).
   */
  private Map<String,Cache> caches = new HashMap<String,Cache>();

  /**
   * There is one lock object per cache.
   */
  private ConcurrentMap<String,Lock> locks = new ConcurrentHashMap<String,Lock>();

  private final KeyGenerator keyGenerator;

  public NonBlockingCachingMethodInterceptor(int cacheTimeSeconds, int firstCacheWaitTimeMillis, int updatedCacheWaitTimeMillis) {
    this(cacheTimeSeconds, firstCacheWaitTimeMillis, updatedCacheWaitTimeMillis, new SimpleKeyGenerator());
  }

  public NonBlockingCachingMethodInterceptor(int cacheTimeSeconds, int firstCacheWaitTimeMillis, int updatedCacheWaitTimeMillis, KeyGenerator keyGenerator) {
    this.cacheTimeSeconds = cacheTimeSeconds;
    this.firstCacheWaitTimeMillis = firstCacheWaitTimeMillis;
    this.updatedCacheWaitTimeMillis = updatedCacheWaitTimeMillis;
    this.keyGenerator = keyGenerator;
  }

  public void setCacheTimeSeconds(int cacheTimeSeconds) {
    this.cacheTimeSeconds = cacheTimeSeconds;
  }

  public void setFirstCacheWaitTimeMillis(int firstCacheWaitTimeMillis) {
    this.firstCacheWaitTimeMillis = firstCacheWaitTimeMillis;
  }

  public void setUpdatedCacheWaitTimeMillis(int updatedCacheWaitTimeMillis) {
    this.updatedCacheWaitTimeMillis = updatedCacheWaitTimeMillis;
  }

  /**
   * This is the method that is invoked over the intercepted method in order
   * to return the cached value wherever possible.
   */
  public Object invoke(MethodInvocation invocation) throws Throwable {

    String methodKey = keyGenerator.generateKey( invocation );

    Cache cache = caches.get(methodKey);

    if( cache == null ) {
      // We haven't cached anything yet.
      logger.debug( "No cache entry for method key {}", methodKey );
      return getUpdatedCache(methodKey, invocation, firstCacheWaitTimeMillis);
    }

    if( new Date().before( cache.getExpiryTime() ) ) {
      // Cache is up to date, we can just return it.
      logger.debug( "Cache hit for method key {}", methodKey );
      return cache.getValue();
    }

    // Cache is out of date.
    logger.debug( "Cache expired for method key {}", methodKey );
    return getUpdatedCache(methodKey, invocation, updatedCacheWaitTimeMillis);

  }

  /**
   * This is used when we need to update the cache, either because we don't have
   * a value cached yet, or the cache has expired.
   *
   * When we haven't got a cached value yet, we have to pull the data out now.
   * All threads must return up to date data, however we can reduce resource
   * contention by only allowing one thread to make the underlying call. The
   * other threads will wait for this thread to complete its call - then it can
   * return the new value. If they wait for the specified amount of time without
   * the lock becoming available though then they give up and just make the call
   * themselves too.
   *
   * When we have an expired cache and we want to get it up to date, again only
   * one thread should update the data. The others should wait on the data being
   * updated if it is within the specified amount of time, otherwise they can
   * simply return the stale cache.
   */
  private Object getUpdatedCache(String methodKey, MethodInvocation invocation, int lockWaitTimeMillis) throws Throwable {

    try {

      Lock newLock = new ReentrantLock();

      // Atomic operation to ensure we use the same lock object for this
      // method as all the other threads.
      Lock oldLock = locks.putIfAbsent(methodKey, newLock);

      Lock lock = oldLock != null ? oldLock : newLock;

      logger.trace( "Trying lock..." );
      if( lock.tryLock( lockWaitTimeMillis, TimeUnit.MILLISECONDS ) ) {
        try {
          // Check the cache again in case it has just been updated by another thread in the time we
          // were waiting for the lock.
          Cache cache = caches.get(methodKey);
          if( cache == null || new Date().after( cache.getExpiryTime() ) ) {
            updateCache(methodKey, invocation);
          }
        } finally {
          logger.trace( "Unlocking" );
          lock.unlock();
        }
      } else {
        logger.trace( "Failed to get lock" );
        // If the cache still doesn't exist then we need to update it anyway. This could happen when the
        // cache is getting created for the first time, if the first thread takes too long to complete.
        // In this case multiple threads could be writing to the cache simultaneously, but we have to let that happen.
        if( caches.get(methodKey) == null ) updateCache(methodKey, invocation);
      }
    } catch( InterruptedException ex ) {
      logger.warn( "Interrupted while waiting for lock for cache with methodKey {}:", methodKey, ex );
    }

    Cache cache = caches.get(methodKey);
    if( cache == null ) {
      // This shouldn't happen, but just in case...
      logger.error( "Cache for method key {} is still null", methodKey );
      return invocation.proceed();
    }

    return cache.getValue();

  }

  protected Cache updateCache(String methodKey, MethodInvocation invocation) throws Throwable {

    logger.debug( "Updating cache for methodKey {}", methodKey );

    // Fall through to the method to get the result

    Date cachedUntil = new Date( System.currentTimeMillis() + (cacheTimeSeconds * 1000) );
    logger.debug( "Cache valid until {} for methodKey {}", cachedUntil, methodKey );

    Cache newCache = new Cache(invocation.proceed(), cachedUntil);
    caches.put(methodKey, newCache);

    return newCache;

  }

  public void expireAllCaches() {

    for( Map.Entry<String,Cache> entry : caches.entrySet() ) {
      Cache cache = entry.getValue();
      cache.setExpiryTime( new Date( System.currentTimeMillis() - 1 ) );
    }

  }


  public class Cache {

    private Object value;
    private Date expiryTime;

    public Cache(Object value, Date expiryTime) {
      this.value = value;
      this.expiryTime = expiryTime;
    }

    public Object getValue() {
      return value;
    }

    public void setValue(Object value) {
      this.value = value;
    }

    public Date getExpiryTime() {
      return expiryTime;
    }

    public void setExpiryTime(Date expiryTime) {
      this.expiryTime = expiryTime;
    }

  }

  public interface KeyGenerator {

    String generateKey(MethodInvocation invocation);

  }

  public static class SimpleKeyGenerator implements KeyGenerator {

    public String generateKey(MethodInvocation invocation) {

      StringBuilder buf = new StringBuilder();

      String targetName = invocation.getThis().getClass().getName();
      String methodName = invocation.getMethod().getName();
      buf.append(targetName).append('.').append(methodName);

      Object[] args = invocation.getArguments();

      buf.append('(');

      for ( int i = 0; i< args.length; i++) {

        if( i > 0 ) buf.append(",");

        buf.append(args[i]);

      }

      buf.append(')');

      return buf.toString();
    }

  }

}
